package org.fujay.commons.langchain4j.config;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.fujay.commons.langchain4j.rag.store.MilvusProperties;
import org.fujay.commons.langchain4j.rag.store.OpenSearchProperties;
import org.opensearch.client.RestClient;
import org.opensearch.client.json.jackson.JacksonJsonpMapper;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.opensearch.client.transport.rest_client.RestClientTransport;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;

import java.util.concurrent.TimeUnit;

/**
 * @author slm
 * @description RAG配置类
 */
@Slf4j
@AutoConfiguration
@RequiredArgsConstructor
@EnableConfigurationProperties({ MilvusProperties.class, OpenSearchProperties.class })
public class RagConfig {

    private final MilvusProperties milvusProperties;
    private final OpenSearchProperties openSearchProperties;

    @Bean
    @ConditionalOnProperty(prefix = "milvus", name = "enable", havingValue = "true")
    public MilvusServiceClient milvusClient() {
        log.info("初始化Milvus客户端连接：host={}, port={}", milvusProperties.getHost(), milvusProperties.getPort());

        ConnectParam.Builder builder = ConnectParam.newBuilder();

        // 优先使用URI，如果没有则使用host和port
        if (milvusProperties.getUri() != null && !milvusProperties.getUri().isEmpty()) {
            builder.withUri(milvusProperties.getUri());
        } else {
            builder.withHost(milvusProperties.getHost())
                    .withPort(milvusProperties.getPort());
        }

        // 设置连接超时时间
        builder.withConnectTimeout(milvusProperties.getConnectTimeoutMs(), TimeUnit.MILLISECONDS);

        // 设置认证信息（如果有）
        if (milvusProperties.getToken() != null && !milvusProperties.getToken().isEmpty()) {
            builder.withToken(milvusProperties.getToken());
        }

        if (milvusProperties.getUsername() != null && !milvusProperties.getUsername().isEmpty()
                && milvusProperties.getPassword() != null && !milvusProperties.getPassword().isEmpty()) {
            builder.withAuthorization(milvusProperties.getUsername(), milvusProperties.getPassword());
        }

        // 创建客户端
        return new MilvusServiceClient(builder.build());
    }

    /**
     * 创建Milvus向量存储
     */
    @Bean
    @ConditionalOnProperty(prefix = "milvus", name = "enable", havingValue = "true")
    public EmbeddingStore<TextSegment> milvusEmbeddingStore(MilvusServiceClient milvusClient) {
        log.info("初始化Milvus向量存储：collection={}, dimension={}",
                milvusProperties.getCollectionName(),
                milvusProperties.getDimension());

        MilvusEmbeddingStore.Builder builder = MilvusEmbeddingStore.builder()
                .dimension(milvusProperties.getDimension())
                .collectionName(milvusProperties.getCollectionName());

        // 使用现有客户端
        try {
            // 反射设置客户端字段 - langchain4j的不同版本API可能不同
            builder.getClass().getMethod("client", MilvusServiceClient.class)
                    .invoke(builder, milvusClient);
        } catch (Exception e) {
            // 如果不支持client方法，则使用host和port
            builder.host(milvusProperties.getHost())
                    .port(milvusProperties.getPort());

            // 设置认证信息
            if (milvusProperties.getUsername() != null && !milvusProperties.getUsername().isEmpty()) {
                builder.username(milvusProperties.getUsername())
                        .password(milvusProperties.getPassword());
            }
        }
        return builder.build();
    }

    /**
     * @author sunliming
     * @description 默认内存向量数据库
     */
    @Bean
    @ConditionalOnMissingBean
    public EmbeddingStore<TextSegment> embeddingStore() {
        return new InMemoryEmbeddingStore<>();
    }

    /**
     * OpenSearch客户端
     */
    @Bean
    public OpenSearchClient openSearchClient() {

        HttpHost host = new HttpHost(openSearchProperties.getOpensearchHost(),
                openSearchProperties.getOpensearchPort());

        final BasicCredentialsProvider credentialsProvider = new BasicCredentialsProvider();
        credentialsProvider.setCredentials(new AuthScope(host), new UsernamePasswordCredentials(
                openSearchProperties.getUserName(), openSearchProperties.getPassword()));

        RestClient restClient = RestClient.builder(host).build();
        /*
         * .
         * setHttpClientConfigCallback(httpClientBuilder ->
         * httpClientBuilder.setDefaultCredentialsProvider(
         * credentialsProvider)).build();
         */

        return new OpenSearchClient(new RestClientTransport(restClient, new JacksonJsonpMapper()));
    }
}
